{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial: CommonRoad Search \n",
    "## Graph Search-Based Motion Planners with Motion Primitives\n",
    "\n",
    "This tutorial shows how [CommonRoad Search](https://gitlab.lrz.de/tum-cps/commonroad-search), or **Graph Search-Based Motion Planner with Motion Primitives**, is used to search for trajectories that connect an **initial state** and a **goal region**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Preparation\n",
    "* Before you proceed, make sure you have read through the tutorial for [CommonRoad-io](https://commonroad.in.tum.de/static/docs/commonroad-io/index.html). Its tutorial can be found [here](https://commonroad.in.tum.de/tutorials/).\n",
    "* Make sure you have installed all necessary modules for CommonRoad Search according to the installation manual.\n",
    "* A documentation for the API of CommonRoad Search can be found [here](https://commonroad.in.tum.de/static/docs/commonroad-search/index.html).\n",
    "* The API of CommonRoad-io can be found [here](https://commonroad.in.tum.de/static/docs/commonroad-io/api/index.html)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Load Python Environment\n",
    "We first import necessary modules for motion planning, as well as some extensions for our convenience."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# enabling autoreload will reload modules automatically before entering the execution of code,\n",
    "# so you could edit the code for your motion planner and execute it right away without reloading again\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "# always show execution time for each cell\n",
    "%load_ext autotime\n",
    "\n",
    "# use notebook to get interactive plots\n",
    "%matplotlib notebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# append main directory\n",
    "import sys\n",
    "sys.path.append(\"../../GSMP/motion_automata\")\n",
    "\n",
    "# load necessary modules and functions\n",
    "from automata.HelperFunctions import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Secondly, we specify the motion planner that we want to use:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Uncomment the following to try out exemplary motion planners\n",
    "\n",
    "# Greedy Best First Search\n",
    "from automata.MotionPlanner_gbfs import MotionPlanner\n",
    "\n",
    "# A* search\n",
    "# from automata.MotionPlanner_Astar import MotionPlanner\n",
    "\n",
    "# Uncomment the following to load your own motion planner\n",
    "\n",
    "# from automata.MotionPlanner import MotionPlanner"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Load Scenarios\n",
    "We call automata.HelperFunctions.load_scenario() to load scenarios and planning problems. Scenarios beginning with letter C stand for cooeprative driving scenarios and possess multiple planning problems. In this case you should provide solutions to all planning problems."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "scenario_path_prefix = '../../scenarios/exercise/'\n",
    "scenario_id = 'USA_Lanker-1_2_T-1'\n",
    "# construct file path\n",
    "scenario_path = scenario_path_prefix + scenario_id + '.xml'\n",
    "\n",
    "scenario, planning_problem_set = load_scenario(scenario_path)\n",
    "\n",
    "# plot the scenario and planning problem set\n",
    "plt.figure(figsize=(8, 8))\n",
    "draw_object(scenario)\n",
    "draw_object(planning_problem_set)\n",
    "plt.gca().set_aspect('equal')\n",
    "plt.margins(0,0)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Load Motion Primitives\n",
    "We call automata.HelperFunctions.generate_automata() to generate a motion automata instance.\n",
    "\n",
    "The variable veh_type_id specifies the id of the vehicle model that we want to use:\n",
    "\n",
    "1: FORD_ESCORT\n",
    "\n",
    "2: BMW320i\n",
    "\n",
    "3: VW_VANAGON"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "veh_type_id = 2\n",
    "\n",
    "if veh_type_id   == 1: veh_type = VehicleType.FORD_ESCORT\n",
    "elif veh_type_id == 2: veh_type = VehicleType.BMW_320i\n",
    "elif veh_type_id == 3: veh_type = VehicleType.VW_VANAGON\n",
    "    \n",
    "automata = generate_automata(veh_type_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Initialize Motion Planner\n",
    "\n",
    "In this step, we would like to configure the planning problem for the motion planner. Each **planning problem** has an individual **initial state** and a **goal region**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# retrieve planning problem with given index (for cooperative scenario:0, 1, 2, ..., otherwise: 0)\n",
    "planning_problem_idx = 0\n",
    "planning_problem = list(planning_problem_set.planning_problem_dict.values())[planning_problem_idx]\n",
    "\n",
    "# add initial state of planning problem to automata\n",
    "automata, initial_motion_primitive = add_initial_state_to_automata(automata, planning_problem)\n",
    "\n",
    "# construct motion planner.\n",
    "motion_planner = MotionPlanner(scenario, planning_problem, automata)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Search for Solution\n",
    "We call automata.HelperFunctions.start_search() to start searching for a feasible solution base on the given motion planner."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result_path, result_dict_status = start_search(scenario,\n",
    "                                               planning_problem, \n",
    "                                               automata,\n",
    "                                               motion_planner, \n",
    "                                               initial_motion_primitive)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Visualize Planned Trajectory\n",
    "Given that we have found a feasible solution, we can also create an interactive visualizor for the solution.\n",
    "We define two helper functions to retrieve state and its collision object at a given time step:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_state_at_time(t):\n",
    "    for state in result_path:\n",
    "        # return the first state that has time_step >= given t\n",
    "        if state.time_step >= t:\n",
    "            return state\n",
    "    # else return last state\n",
    "    return path[-1]\n",
    "\n",
    "def draw_state(t):\n",
    "    print(\"current time step: \", t)\n",
    "    plt.clf()\n",
    "    draw_object(scenario, draw_params={'time_begin': t})\n",
    "    draw_object(planning_problem)\n",
    "    plt.gca().set_aspect('equal')\n",
    "    plt.gca().set_axis_off()\n",
    "    plt.margins(0,0)\n",
    "    plt.show()\n",
    "\n",
    "    if result_path is not None:\n",
    "        state = get_state_at_time(t)\n",
    "        trajectory = Trajectory(initial_time_step=int(state.time_step),state_list=[state])\n",
    "        prediction = TrajectoryPrediction(trajectory=trajectory, shape=automata.egoShape)\n",
    "        collision_object = create_collision_object(prediction)\n",
    "        draw_it(collision_object)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(8,8))\n",
    "\n",
    "# use the slider to have at look at every time step of the planned trajectory\n",
    "widgets.interact(draw_state, t=widgets.IntSlider(min=0,max=result_path[-1].time_step,step=1,value=0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# if you wish to visualize the final trajectory\n",
    "if len(result_dict_status['path_current']) > 0:\n",
    "    plt.figure(figsize=(8, 8))\n",
    "    plt.clf()\n",
    "    draw_object(scenario)\n",
    "    draw_object(planning_problem)\n",
    "    plt.gca().set_aspect('equal')\n",
    "    plt.gca().set_axis_off()\n",
    "    plt.margins(0,0)\n",
    "    plt.show()\n",
    "    \n",
    "    trajectory = Trajectory(0, result_dict_status['path_current'][0:])\n",
    "    prediction = TrajectoryPrediction(trajectory=trajectory, shape=automata.egoShape)\n",
    "    collision_object = create_collision_object(prediction)\n",
    "    draw_it(collision_object, draw_params={'collision': {'facecolor': 'green'}})\n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# if you wish to print out individual state\n",
    "# for state in result_path:\n",
    "#     print(state)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Write Solution to CommonRoad XML File\n",
    "The final step is writing your solution to the planning problem in the scenario into an XML file which is uploadable to our [Benchmark](https://commonroad.in.tum.de/submissions/create).\n",
    "You may also refer to [Solution Writer](https://commonroad.in.tum.de/static/docs/commonroad-io/api/common.html#module-commonroad.common.solution_writer), [Cost Function](https://gitlab.lrz.de/tum-cps/commonroad-cost-functions/blob/master/costFunctions_commonRoad.pdf) and [Vehicle Model](https://gitlab.lrz.de/tum-cps/commonroad-vehicle-models/blob/master/vehicleModels_commonRoad.pdf) for additional information.\n",
    "Note: In this tutorial we use Kinematic Single-Track Model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "list_state = list()\n",
    "\n",
    "for state in result_path:\n",
    "    kwarg = {'position': state.position, \n",
    "             'velocity': state.velocity,\n",
    "             'steering_angle': state.steering_angle, \n",
    "             'orientation': state.orientation, \n",
    "             'time_step': state.time_step}\n",
    "    list_state.append(State(**kwarg))\n",
    "\n",
    "trajectory = Trajectory(initial_time_step=list_state[0].time_step, state_list=list_state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from commonroad.common.solution_writer import CommonRoadSolutionWriter, VehicleModel, VehicleType, CostFunction\n",
    "output_dir=\"../../solutions/\"\n",
    "\n",
    "# write solution to a xml file\n",
    "csw = CommonRoadSolutionWriter(output_dir=output_dir,                   # output directory\n",
    "                               scenario_id=scenario_id,                 # scenario id\n",
    "                               step_size=0.1,                           # time step of scenario, typically 0.1\n",
    "                               vehicle_type=veh_type,                   # vehicle type, change if needed\n",
    "                               vehicle_model=VehicleModel.KS,           # vehicle model, change if needed\n",
    "                               cost_function=CostFunction.SM1)          # cost funtion, ** do not use JB1 **\n",
    "\n",
    "# add trajectory solution\n",
    "csw.add_solution_trajectory(trajectory=trajectory, planning_problem_id=planning_problem.planning_problem_id)\n",
    "\n",
    "csw.write_to_file(overwrite=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (cr37)",
   "language": "python",
   "name": "cr37"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
